import os
import sys

sys.path.append('./')

import cv2
import numpy as np
import torch
import torchvision.transforms as transforms
from models.model_stages import BiSeNet
from steel import Steel
from torch.utils.data import DataLoader
from tqdm import tqdm


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg'):
                img_list.append(os.path.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list

def get_center(contours_pred):
    points_center = []
    for contour in contours_pred:
        x_mean = np.mean(contour[..., 0]) - 10
        y_mean = np.mean(contour[..., 1]) - 10
        points_center.append([x_mean, y_mean])
    
    points_center = np.array(points_center, np.int64)
    keep = np.argsort(points_center[:, 0])
    points_center = points_center[keep]

    return points_center

def eval1(net, img_path, save_dir, inputsize, show=False):
    dsval = Steel(img_path, mode='test', use_heatmap=False, inputsize=inputsize, return_path=True)
    dl = DataLoader(dsval,
                    batch_size = 1,
                    shuffle = False,
                    # sampler = sampler_val,
                    num_workers = 0,
                    drop_last = False)

    ## predict
    thresh = 0.5
    diter = enumerate(tqdm(dl))

    pred_total = 0
    pred_true = 0
    for i, (imgs, label, im_path) in diter:
        imgs = imgs.cuda()
        label = (label.squeeze(1).squeeze(0).numpy() * 255).astype(np.uint8)

        logits = net(imgs)[0]

        probs = torch.sigmoid(logits)
        preds = (probs[0][0] > thresh).cpu().numpy()
        preds = preds.astype(np.uint8) * 255

        # label nums
        contours_label = cv2.findContours(label, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        nums_label = len(contours_label)

        # pred nums
        # 腐蚀处理
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))
        # kernel = np.ones((5, 5), dtype=np.uint8)
        # GaussianKernel = cv2.getGaussianKernel(5, 0, cv2.CV_32F)
        # kernel = GaussianKernel * GaussianKernel.T
        erosion_preds = cv2.erode(preds, kernel, anchor=(-1,-1), iterations=2)
        contours_pred = cv2.findContours(erosion_preds, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        nums_pred = len(contours_pred)

        pred_total += 1
        if nums_label == nums_pred:
            pred_true += 1
        else:
            if show:
                # 获取预测中心点
                points_center = get_center(contours_pred)

                # 画mask
                img = (imgs.cpu().numpy()[0].transpose(1,2,0) * 255).astype(np.uint8)
                img[..., 0][preds==255] = np.array(img[..., 0][preds==255] * 0.5 + 255.0 * 0.5, np.uint8)

                # 画数字
                for idx in range(nums_pred):
                    point = points_center[idx]
                    img = cv2.putText(img.copy(), str(idx+1), (point[0], point[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

                # 统计
                img = cv2.putText(img.copy(), 'nums_label: {}'.format(nums_label), (10,20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                img = cv2.putText(img.copy(), 'nums_pred: {}'.format(nums_pred), (10,50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

                # 保存图片
                img_name = os.path.basename(im_path[0]).split('.')[0]
                
                save_path = os.path.join(save_dir, img_name)
                os.makedirs(save_path, exist_ok=True)
                cv2.imwrite(os.path.join(save_path, "vis_image.png"), img)
                cv2.imwrite(os.path.join(save_path, "preds_mask.png"), preds)
                cv2.imwrite(os.path.join(save_path, "erosion_mask.png"), erosion_preds)
                
                # cv2.imshow('img', img)
                # cv2.imshow('mask', preds)
                # cv2.waitKey(0)

    print("precision: {:.2f}".format(pred_true / pred_total))
    print("error_nums: {}".format(pred_total - pred_true))

def eval2(net, img_path, save_dir, inputsize, show=False):
    dsval = Steel(img_path, mode='test', use_heatmap=False, inputsize=inputsize, return_path=True)
    dl = DataLoader(dsval,
                    batch_size = 1,
                    shuffle = False,
                    # sampler = sampler_val,
                    num_workers = 0,
                    drop_last = False)

    ## predict
    thresh = 0.5
    diter = enumerate(tqdm(dl))

    pred_total = 0
    pred_true = 0
    for i, (imgs, label, im_path) in diter:

        imgs = imgs.cuda()
        label = (label.squeeze(1).squeeze(0).numpy() * 255).astype(np.uint8)

        logits = net(imgs)[0]

        probs = torch.sigmoid(logits)
        preds = (probs[0][0] > thresh).cpu().numpy()
        preds = preds.astype(np.uint8) * 255

        # label nums
        contours_label = cv2.findContours(label, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        nums_label = len(contours_label)

        # noise removal
        kernel = np.ones((3,3),np.uint8)
        opening = cv2.morphologyEx(preds, cv2.MORPH_OPEN,kernel, iterations = 5)
        # cv2.imshow('opening', opening)

        # sure background area
        sure_bg = cv2.dilate(opening,kernel,iterations=1)
        # cv2.imshow('sure_bg', sure_bg)

        # Finding sure foreground area
        dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 3)
        # dist_transform_show = (dist_transform / dist_transform.max() * 255).astype(np.uint8)
        # cv2.imshow('dist_transform_show', dist_transform_show)

        ret, sure_fg = cv2.threshold(dist_transform, 0.2 * dist_transform.max(), 255, 0)
        # cv2.imshow('sure_fg', sure_fg)

        # Finding unknown region
        sure_fg = np.uint8(sure_fg)
        unknown = cv2.subtract(sure_bg,sure_fg)
        # cv2.imshow('unknown', unknown)

        # Marker labelling
        ret, markers = cv2.connectedComponents(sure_fg)
        
        # Add one to all labels so that sure background is not 0, but 1
        markers = markers+1

        # Now, mark the region of unknown with zero
        markers[unknown==255] = 0

        img = (imgs.cpu().numpy()[0].transpose(1,2,0) * 255).astype(np.uint8)
        markers = cv2.watershed(img, markers)

        bg_mask = np.zeros_like(img)
        bg_mask[markers == -1] = [0,0,255]
        # cv2.imshow('bg_mask', bg_mask)

        pix_unique = np.unique(markers)
        area_list = []
        for pix in pix_unique:
            if pix == -1:
                continue
            area = (markers == pix).sum()
            if area > 1600*1600*0.25:
                continue
            area_list.append(area)

        nums_pred = 0
        area_mean = np.mean(area_list)
        for area in area_list:
            if area < area_mean * 0.25:
                continue
            nums_pred += 1
            scale = area / area_mean
            if 1.6 < scale < 2.6:
                nums_pred += 1
            elif scale >= 2.4:
                nums_pred += 2

        pred_total += 1
        if nums_label == nums_pred:
            pred_true += 1
        else:
            if show:
                # 保存图片
                img_name = os.path.basename(im_path[0]).split('.')[0]

                for pix in np.unique(markers):
                    if pix == -1:
                        continue
                    b = np.random.randint(0, 255)
                    g = np.random.randint(0, 255)
                    r = np.random.randint(0, 255)
                    img[markers == pix] = [b, g, r]
                
                # 统计
                img = cv2.putText(img.copy(), 'nums_label: {}'.format(nums_label), (10,20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
                img = cv2.putText(img.copy(), 'nums_pred: {}'.format(nums_pred), (10,50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)


                save_path = os.path.join(save_dir, img_name)
                os.makedirs(save_path, exist_ok=True)
                cv2.imwrite(os.path.join(save_path, "vis_image.png"), img)
                cv2.imwrite(os.path.join(save_path, "preds_mask.png"), preds)
                cv2.imwrite(os.path.join(save_path, "label.png"), label)
                
                # cv2.imshow('img', img)
                # cv2.imshow('mask', preds)
                # cv2.waitKey(0)

    print("precision: {:.2f}".format(pred_true / pred_total))
    print("error_nums: {}".format(pred_total - pred_true))

def infer1(net, img_path, inputsize):
    img_list = get_img_list(img_path)
    to_tensor = transforms.ToTensor()

    for path_i in img_list:
        img_ori = cv2.imread(path_i)
        img = cv2.resize(img_ori, inputsize, interpolation=cv2.INTER_NEAREST)
        
        inp = to_tensor(img)[None, ]
        inp = inp.cuda()

        logits = net(inp)[0]

        probs = torch.sigmoid(logits)
        preds = (probs[0][0] > 0.5).cpu().numpy()
        preds = preds.astype(np.uint8) * 255

        # pred nums
        # 腐蚀处理
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5,5))
        erosion_preds = cv2.erode(preds, kernel, anchor=(-1,-1), iterations=2)
        contours_pred = cv2.findContours(erosion_preds, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
        nums_pred = len(contours_pred)

        # 获取预测中心点
        points_center = get_center(contours_pred)

        # 画mask
        # img = (img.cpu().numpy()[0].transpose(1,2,0) * 255).astype(np.uint8)
        img[..., 0][preds==255] = np.array(img[..., 0][preds==255] * 0.5 + 255.0 * 0.5, np.uint8)

        # 画数字
        for idx in range(nums_pred):
            point = points_center[idx]
            img = cv2.putText(img.copy(), str(idx+1), (point[0], point[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        # 统计
        img = cv2.putText(img.copy(), 'nums_pred: {}'.format(nums_pred), (10,50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)

        # 保存图片
        # cv2.imwrite(osp.join(save_dir, "{}_vis_image.png".format(i)), img)
        cv2.imshow('img', img)
        cv2.waitKey(0)

if __name__ == '__main__':
    eval_mode = True
    respth = r'checkpoints\train_STDC1-Seg\train_0521\pths\model_final.pth'
    save_dir = r'evaluation_logs\v4_res_post2'
    img_path = 'data/shushu/test'

    backbone = 'STDCNet813'
    use_heatmap = False
    inputsize = [1600, 1600]
    use_boundary_2 = False
    use_boundary_4 = False
    use_boundary_8 = False
    use_boundary_16 = False
    use_conv_last = False

    n_classes = 1
    net = BiSeNet(backbone=backbone, n_classes=n_classes,
        use_boundary_2=use_boundary_2, use_boundary_4=use_boundary_4, 
        use_boundary_8=use_boundary_8, use_boundary_16=use_boundary_16, 
        use_conv_last=use_conv_last)
    net.load_state_dict(torch.load(respth))
    net.cuda()
    net.eval()
    
    if eval_mode:
        eval2(net, img_path, save_dir, inputsize, show=True)
    else:
        infer1(net, img_path, inputsize)